import time

import numpy as np
import torch

from tools import GetConfusionMatrix, PrintLog, WriteArray


def train(model, epochs, train_loader, test_loader, optimizer, index):
    # loss_list = []
    param_list = []
    max_T_acc = 0
    with_N_acc = 0
    epsilon = 1e-6
    loss_fn = torch.nn.BCELoss()
    path_name = "pth/red_predict_" + str(index) + ".pth"
    param_name = 'param/param_' + str(index) + '.txt'

    for epoch in range(epochs):
        print('Starting training....{}/{}'.format(epoch + 1, epochs))
        loss_sub_list = []
        start = time.time()
        for data_lists, target_lists in train_loader:
            model.train()
            pred_lists = model(data_lists)
            train_loss = loss_fn(pred_lists, target_lists)
            tmp_loss = train_loss.detach().numpy()
            loss_sub_list.append(tmp_loss)

            optimizer.zero_grad()
            train_loss.backward()
            optimizer.step()

        TP = 0
        FP = 0
        TN = 0
        FN = 0
        for data_lists, target_lists in test_loader:
            model.eval()
            pred_lists = model(data_lists)
            pred_target = torch.max(pred_lists, 1)[1]
            true_target = torch.max(target_lists, 1)[1]
            tmp_TP, tmp_FP, tmp_TN, tmp_FN = GetConfusionMatrix(pred_target, true_target)
            TP += tmp_TP
            FP += tmp_FP
            TN += tmp_TN
            FN += tmp_FN

        end = time.time()
        epoch_loss = np.mean(loss_sub_list)
        # loss_list.append(epoch_loss)
        T_acc = TP / (TP + FN + epsilon)
        N_acc = TN / (TN + FP + epsilon)
        print('Epoch loss: {:f}, T_acc: {:f}, N_acc: {:f}, time used: ({:.1f}s)'.format(epoch_loss, T_acc, N_acc,
                                                                                        end - start))

        if T_acc > max_T_acc:
            torch.save(model.state_dict(), path_name)
            max_T_acc = T_acc
            with_N_acc = N_acc

    print(f"train model {path_name}'s T_acc is {max_T_acc}, N_acc is {with_N_acc}")
    param_list.append(max_T_acc)
    param_list.append(with_N_acc)
    param_list = [param_list]
    WriteArray(param_list, param_name)
